#include "connected_layer.h"
#include "utils.h"
#include "owl_math.h"
#include "network_state.h"
#include "logging.h"
#include "activations.h"

ConnectedLayer MakeConnectedLayer(int batch, int inputs, int outputs, Activation activation, int batch_normalize) {
    int i = 0;
    ConnectedLayer layer;// = {0};
    layer.type = CONNECTED;
    layer.inputs = inputs;
    layer.outputs = outputs;
    layer.batch = batch;
    layer.batch_normalize = batch_normalize;

    layer.h = 1;
    layer.w = 1;
    layer.c = inputs;
    layer.out_h = 1;
    layer.out_w = 1;
    layer.out_c = outputs;

    layer.data_output = (Float*)OwlAlloc(batch * outputs, sizeof(Float));
    layer.data_delta = (Float*)OwlAlloc(batch * outputs, sizeof(Float));

    layer.data_weights = (Float*)OwlAlloc(outputs * inputs, sizeof(Float));
    layer.data_biases = (Float*)OwlAlloc(outputs, sizeof(Float));
    layer.data_weight_updates = (Float*)OwlAlloc(inputs * outputs, sizeof(Float));
    layer.data_bias_updates = (Float*)OwlAlloc(outputs, sizeof(Float));

    layer.Forward = ForwardConnectedLayer;
    layer.Backward = BackwardConnecteLayer;
    layer.Update = UpdateConnectedLayer;

    layer.activation = activation;

    Float scale = sqrt(2.0 / inputs);
    for (i = 0; i < outputs * inputs; ++i) {
        layer.data_weights[i] = i; //scale * rand_uniform(-1, 1);
    }
    for (i = 0; i < outputs; ++i) {
        layer.data_biases[i] = 0.0;
    }

    if (batch_normalize) {
        layer.data_scales = (Float*)OwlAlloc(outputs, sizeof(Float));
        layer.data_scale_updates = (Float*)OwlAlloc(outputs, sizeof(Float));
        for (i = 0; i < outputs; ++i) {
            layer.data_scales[i] = 1;
        }

        layer.data_mean = (Float*)OwlAlloc(outputs, sizeof(Float));
        layer.data_mean_delta = (Float*)OwlAlloc(outputs, sizeof(Float));
        layer.data_variance = (Float*)OwlAlloc(outputs, sizeof(Float));
        layer.data_variance_delta = (Float*)OwlAlloc(outputs, sizeof(Float));

        layer.data_rolling_mean = (Float*)OwlAlloc(outputs, sizeof(Float));
        layer.data_rolling_variance = (Float*)OwlAlloc(outputs, sizeof(Float));

        layer.data_x = (Float*)OwlAlloc(batch*outputs, sizeof(Float));
        layer.data_x_norm = (Float*)OwlAlloc(batch * outputs, sizeof(Float));
    }

    return layer;
}

void ForwardConnectedLayer(ConnectedLayer layer, NetworkState state) {
    LOGFUNC("ForwardConnectedLayer");
    int i;
    OwlFill(layer.outputs * layer.batch, 0, layer.data_output, 1);
    int m = layer.batch;
    int k = layer.inputs;
    int n = layer.outputs;
    Float *a = state.data_input;
    OWLLOG("ConnectedLayer data_input ", a, k);
    Float *b = layer.data_weights;
    OWLLOG("ForwardConnectedLayer weights", layer.data_weights, 6);
    Float *c = layer.data_output;
    OwlGemm(0, 1, m, n, k, 1, a, k, b, k, 1, c, n);
    OWLLOG("ForwardConnectedLayer data output ", layer.data_output, 3);
    // if (layer.batch_normalize) {
    //     if (state.train) {
    //         OwlMean(layer.data_output, layer.batch, layer.outputs, 1, layer.data_mean);
    //         OwlVariance(layer.data_output, layer.data_mean, layer.batch, layer.outputs, 1, layer.data_variance);
    //         OwlScal(layer.outputs, 0.95, layer.data_rolling_mean, 1);
    //         OwlAxpy(layer.outputs, 0.05, layer.data_mean, 1, layer.data_rolling_mean, 1);
    //         OwlScal(layer.outputs, 0.95, layer.data_rolling_variance, 1);
    //         OwlAxpy(layer.outputs, 0.05, layer.data_variance, 1, layer.data_rolling_variance, 1);
    //
    //         OwlCopy(layer.outputs * layer.batch, layer.data_output, 1, layer.data_x, 1);
    //         OwlNormalize(layer.data_output, layer.data_mean, layer.data_variance, layer.batch, layer.outputs, 1);
    //         OwlCopy(layer.outputs * layer.batch, layer.data_output, 1, layer.data_x_norm, 1);
    //     } else {
    //         OwlNormalize(layer.data_output, layer.data_rolling_mean, layer.data_rolling_variance, layer.batch, layer.outputs, 1);
    //     }
    //
    //     //OwlScaleBias(layer.output, layer.scales, layer.batch, layer.outputs, 1);
    // }
    for (i = 0; i < layer.batch; ++i) {
        OwlAxpy(layer.outputs, 1, layer.data_biases, 1, layer.data_output + i * layer.outputs, 1);
    }
    ActivateArray(layer.outputs * layer.batch, layer.data_output, layer.activation);
    OWLLOG("ForwardConnectedLayer", layer.data_output, 6);
}

void BackwardConnecteLayer(ConnectedLayer layer, NetworkState state) {
    GradientArray(layer.batch * layer.outputs, layer.data_output, layer.activation, layer.data_delta);
    int i = 0;
    for (i = 0; i < layer.batch; ++i) {
        OwlAxpy(layer.outputs, 1, layer.data_delta + i * layer.outputs, 1, layer.data_bias_updates, 1);
    }
    // if (layer.batch_normalize) {
    //
    // }
    OWLLOG("BackwardConnecteLayer biase ", layer.data_bias_updates, 1);
    int m = layer.outputs;
    int k = layer.batch;
    int n = layer.inputs;
    Float *A = layer.data_delta;
    Float *B = state.data_input;
    Float *C = layer.data_weight_updates;
    OwlGemm(1, 0, m, n, k, 1, A, m, B, n, 1, C, n);
    OWLLOG("BackwardConnecteLayer input ", B, k * n);
    OWLLOG("BackwardConnecteLayer weights ", C, m * n);
    m = layer.batch;
    k = layer.outputs;
    n = layer.inputs;
    A = layer.data_delta;
    B = layer.data_weights;
    C = state.data_delta;
    //OWLLOG("BackwardConnecteLayer, data_delta", C, 1);
    if (C) {
        OwlGemm(0, 0, m, n, k, 1, A, k, B, n, 1, C, n);
    }
    //OWLLOG("BackwardConnecteLayer, data_delta", C, 1);
}
void UpdateConnectedLayer(ConnectedLayer layer, int batch, float learning_rate, float momentum, float decay) {
    OwlAxpy(layer.outputs, learning_rate / batch, layer.data_bias_updates, 1, layer.data_biases, 1);
    OwlScal(layer.outputs, momentum, layer.data_bias_updates, 1);
    // if (layer.batch_normalize) {
    //     OwlAxpy(layer.outputs, learning_rate / batch, layer.data_scale_updates, 1,
    //     layer.data_scales, 1);
    //     OwlScal(layer.outputs, momentum, layer.data_scale_updates, 1);
    // }

    OwlAxpy(layer.inputs * layer.outputs, -decay * batch, layer.data_weights, 1,
    layer.data_weight_updates, 1);
    OwlAxpy(layer.inputs * layer.outputs, learning_rate / batch, layer.data_weight_updates, 1, layer.data_weights, 1);
    OwlScal(layer.inputs * layer.outputs, momentum, layer.data_weight_updates, 1);
}
